#!/usr/bin/env python
# -*- conding: utf-8 -*-

"""
@Time     : 2024/10/25 6:34
@Author   : liujingmao
@File     : 2.并行节点.py
"""

from typing import Any

from langchain_core.messages import AIMessage, HumanMessage
from langgraph.graph.message import StateGraph, MessagesState

graph_builder = StateGraph(MessagesState)


def chatbot(state: MessagesState, config: dict) -> Any:
    return {"messages": [AIMessage(content="你好，我是OpenAI开发的聊天机器人")]}


def parallel1(state: MessagesState, config: dict) -> Any:
    print("并行1: ", state)
    return {"messages": [HumanMessage(content="这是并行1函数")]}


def parallel2(state: MessagesState, config: dict) -> Any:
    print("并行2: ", state)
    return {"messages": [HumanMessage(content="这是并行2函数")]}


def chat_end(state: MessagesState, config: dict) -> Any:
    print("聊天结束: ", state)
    return {"messages": [HumanMessage(content="这是聊天结束函数")]}


graph_builder.add_node("chat_bot", chatbot)
graph_builder.add_node("parallel1", parallel1)
graph_builder.add_node("parallel2", parallel2)
graph_builder.add_node("chat_end", chat_end)

graph_builder.set_entry_point("chat_bot")
graph_builder.set_finish_point("chat_end")
graph_builder.add_edge("chat_bot", "parallel1")
graph_builder.add_edge("chat_bot", "parallel2")
graph_builder.add_edge("parallel2", "chat_end")

graph = graph_builder.compile()

print(graph.invoke({"messages": [HumanMessage(content="你好，你是")]}))

"""
环流程实现工具调用Agent\2.并行节点.py 
并行1:  {'messages': [HumanMessage(content='你好，你是', additional_kwargs={}, response_metadata={}, id='994ba9ea-2722-442a-a2b1-51b04cc03fb3'), AIMessage(content='你好，我是OpenAI开发的聊天机器人', additional_kwargs={}, response_metadata={}, id='cf089dbd-7381-456f-8a34-abc8cc7a2a0c')]}
并行2:  {'messages': [HumanMessage(content='你好，你是', additional_kwargs={}, response_metadata={}, id='994ba9ea-2722-442a-a2b1-51b04cc03fb3'), AIMessage(content='你好，我是OpenAI开发的聊天机器人', additional_kwargs={}, response_metadata={}, id='cf089dbd-7381-456f-8a34-abc8cc7a2a0c')]}
聊天结束:  {'messages': [HumanMessage(content='你好，你是', additional_kwargs={}, response_metadata={}, id='994ba9ea-2722-442a-a2b1-51b04cc03fb3'), AIMessage(content='你好，我是OpenAI开发的聊天机器人', additional_kwargs={}, response_metadata={}, id='cf089dbd-7381-456f-8a34-abc8cc7a2a0c'), HumanMessage(content='这是并行1函数', additional_kwargs={}, response_metadata={}, id='c65176bb-7812-4fed-b19a-bfa90eca8c68'), HumanMessage(content='这是并行2函数', additional_kwargs={}, response_metadata={}, id='07780a3f-f0d1-4db9-a58a-a4ae0fc4c6da')]}
{'messages': [HumanMessage(content='你好，你是', additional_kwargs={}, response_metadata={}, id='994ba9ea-2722-442a-a2b1-51b04cc03fb3'), AIMessage(content='你好，我是OpenAI开发的聊天机器人', additional_kwargs={}, response_metadata={}, id='cf089dbd-7381-456f-8a34-abc8cc7a2a0c'), HumanMessage(content='这是并行1函数', additional_kwargs={}, response_metadata={}, id='c65176bb-7812-4fed-b19a-bfa90eca8c68'), HumanMessage(content='这是并行2函数', additional_kwargs={}, response_metadata={}, id='07780a3f-f0d1-4db9-a58a-a4ae0fc4c6da'), HumanMessage(content='这是聊天结束函数', additional_kwargs={}, response_metadata={}, id='0bb3e92a-140b-4e7f-bcc5-b86caaf71ba9')]}
"""
